import torchtext
import torch
import torch.nn as nn
from pytorch深度学习应用实战.第12章.词汇处理 import create_vocabulary

maxlen = 4      # 语句最大字数
# 测试资料
docs = ['Well done!',
        'Good work',
        'Great effort',
        'nice work',
        'Excellent!',
        'Weak',
        'Poor effort!',
        'not good',
        'poor work',
        'Could have done better']

vocab_object, clean_text_list, clean_index_list = create_vocabulary(docs)

# 若字串过长，删除多余单字
clean_index_list = torchtext.functional.truncate(clean_index_list, maxlen)

# 若字串长度不足，后面补 0
while len(clean_index_list[0]) < maxlen:
    clean_index_list[0] += [0]
torchtext.functional.to_tensor(clean_index_list, 0) # 0:不足补0

embeds = nn.Embedding(vocab_object.__len__(), 5)
X = torchtext.functional.to_tensor(clean_index_list, 0) # 0:不足补0
embed_output = embeds(X)
print(embed_output.shape)


class RecurrentNet(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.fc = nn.Linear(embed_dim * maxlen, num_class) # 要乘以 maxlen
        self.embed_dim = embed_dim
        self.init_weights()

    def init_weights(self):
        initrange = 0.5
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()

    def forward(self, text):
        embedded = self.embedding(text)
        out = embedded.reshape(embedded.size(0), -1) # 转换成1维
        return self.fc(out)

model = RecurrentNet(vocab_object.__len__(), 10, 1)

# 定义 10 个语句的正面(1)或负面(0)的情绪
y = torch.FloatTensor([1,1,1,1,1,0,0,0,0,0])
X = torchtext.functional.to_tensor(clean_index_list, 0) # 0:不足补0

# 指定优化器、损失函数
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())

# 模型训练
for epoch in range(1000):
    outputs = model.forward(X) #forward pass
    optimizer.zero_grad()
    loss = criterion(outputs.reshape(-1), y)
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        #print(outputs.shape)
        print("Epoch: {epoch}, loss: {loss.item():1.5f}")